{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "9CagYlhclDR4"
   },
   "source": [
    "# Ransac and Outlier Removal\n",
    "\n",
    "## Notebook Setup \n",
    "The following cell will install Drake, checkout the manipulation repository, and set up the path (only if necessary).\n",
    "- On Google's Colaboratory, this **will take approximately two minutes** on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  \n",
    "\n",
    "More details are available [here](http://manipulation.mit.edu/drake.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "3kqzbo_AlDR6"
   },
   "outputs": [],
   "source": [
    "import importlib\n",
    "import os, sys\n",
    "from urllib.request import urlretrieve\n",
    "\n",
    "if 'google.colab' in sys.modules and importlib.util.find_spec('manipulation') is None:\n",
    "    urlretrieve(f\"http://manipulation.csail.mit.edu/scripts/setup/setup_manipulation_colab.py\",\n",
    "                \"setup_manipulation_colab.py\")\n",
    "    from setup_manipulation_colab import setup_manipulation\n",
    "    setup_manipulation(manipulation_sha='db19656799c11a58bc37dd20604ad38eafb09c62', drake_version='20200921', drake_build='nightly')\n",
    "\n",
    "# python libraries\n",
    "import numpy as np\n",
    "\n",
    "# Determine if this notebook is currently running as a notebook or a unit test.\n",
    "from IPython import get_ipython\n",
    "running_as_notebook = get_ipython() and hasattr(get_ipython(), 'kernel')\n",
    "\n",
    "# Start a single meshcat server instance to use for the remainder of this notebook.\n",
    "from meshcat.servers.zmqserver import start_zmq_server_as_subprocess\n",
    "server_args = []\n",
    "if 'google.colab' in sys.modules:\n",
    "    server_args = ['--ngrok_http_tunnel']\n",
    "proc, zmq_url, web_url = start_zmq_server_as_subprocess(server_args=server_args)\n",
    "\n",
    "# TODO(russt): upstream this to drake\n",
    "import meshcat.geometry as g\n",
    "import meshcat.transformations as tf\n",
    "\n",
    "from pydrake.all import RigidTransform, RotationMatrix, RollPitchYaw\n",
    "\n",
    "import open3d as o3d \n",
    "import meshcat\n",
    "\n",
    "from IPython.display import clear_output\n",
    "clear_output()\n",
    "\n",
    "from manipulation import FindResource\n",
    "\n",
    "# Visualize Stanford Bunny \n",
    "pcd = o3d.io.read_point_cloud(FindResource(\"models/bunny/bun_zipper_res2.ply\"))\n",
    "pointcloud_model = np.asarray(pcd.points).transpose()\n",
    "\n",
    "# First, clean the origin a bit to define nominal pose. \n",
    "X = np.array([[1., 0., 0., 0.0],\n",
    "              [0., np.cos(np.pi/2), -np.sin(np.pi/2), 0.],\n",
    "              [0., np.sin(np.pi/2), np.cos(np.pi/2), -0.05]])\n",
    "Xtemp = RigidTransform(X)\n",
    "X = np.array([[np.cos(np.pi/2), -np.sin(np.pi/2), 0, 0.],\n",
    "              [np.sin(np.pi/2), np.cos(np.pi/2), 0., 0.],\n",
    "              [0., 0., 1., 0.]])\n",
    "X = RigidTransform(X).multiply(Xtemp)\n",
    "\n",
    "pointcloud_model = X.multiply(pointcloud_model)\n",
    "\n",
    "# point clouds of planar surface\n",
    "import numpy as np\n",
    "grid_spec = 50\n",
    "xy_axis = np.linspace(-0.5, 0.5, grid_spec)\n",
    "plane_x, plane_y = np.meshgrid(xy_axis, xy_axis)\n",
    "points_plane_xy = np.c_[plane_x.flatten(), plane_y.flatten(), np.zeros(grid_spec**2)]\n",
    "bunny_w_plane = np.c_[points_plane_xy.T, pointcloud_model]\n",
    "\n",
    "def fit_plane(xyzs):\n",
    "    '''\n",
    "    Args:\n",
    "      xyzs is (N, 3) numpy array\n",
    "    Returns:\n",
    "      (4,) numpy array\n",
    "    '''\n",
    "    center = np.mean(xyzs, axis=0)\n",
    "    cxyzs = xyzs - center\n",
    "    U, S, V = np.linalg.svd(cxyzs)\n",
    "    normal = V[-1]              # last row of V\n",
    "    d = -center.dot(normal)\n",
    "    plane_equation = np.hstack([normal, d])\n",
    "    return plane_equation\n",
    "\n",
    "# visualize a facet\n",
    "def DrawFacet(vis, abcd, name, center=None,\n",
    "              prefix='facets', radius=0.02, thickness=0.001, color=0xffffff, opacity=0.6):\n",
    "    normal = np.array(abcd[:3]).astype(float)\n",
    "    normal /= np.linalg.norm(normal)\n",
    "    d = -abcd[3] / np.linalg.norm(normal)\n",
    "\n",
    "    R = np.eye(3)\n",
    "    R[:, 2] = normal\n",
    "    z = normal\n",
    "    if abs(z[0]) < 1e-8:\n",
    "        x = np.array([0, -normal[2], normal[1]])\n",
    "    else:\n",
    "        x = np.array([-normal[1], normal[0], 0])\n",
    "    x /= np.linalg.norm(x)\n",
    "    R[:, 0] = x\n",
    "    R[:, 1] = np.cross(z, x)\n",
    "\n",
    "    X = np.eye(4)\n",
    "    Rz = RollPitchYaw(np.pi/2, 0, 0).ToRotationMatrix().matrix()\n",
    "    X[:3, :3] = R.dot(Rz)\n",
    "    if center is None:\n",
    "        X[:3, 3] = d * normal\n",
    "    else:\n",
    "        X[:3, 3] = center\n",
    "              \n",
    "    X_normal = X.copy()\n",
    "    X_normal[:3, :3] = R\n",
    "    \n",
    "    material = meshcat.geometry.MeshLambertMaterial(\n",
    "        color=color, opacity=opacity)\n",
    "    \n",
    "    vis[prefix][name][\"plane\"].set_object(\n",
    "        meshcat.geometry.Cylinder(thickness, radius), material)\n",
    "    \n",
    "    normal_vertices = np.array([[0, 0, 0], [0, 0, radius]]).astype(float)\n",
    "    vis[prefix][name][\"normal\"].set_object(\n",
    "        meshcat.geometry.Line(meshcat.geometry.PointsGeometry(normal_vertices.T)))\n",
    "    \n",
    "    vis[prefix][name][\"plane\"].set_transform(X)\n",
    "    vis[prefix][name][\"normal\"].set_transform(X_normal)\n",
    "\n",
    "def generate_color_mat(color_vec, shape):\n",
    "  color_mat = np.tile(np.array(color_vec).astype(np.float32).reshape(3,1), (1, shape[1]))\n",
    "  return color_mat\n",
    "\n",
    "vis = meshcat.Visualizer(zmq_url)\n",
    "# vis = meshcat.Visualizer(zmq_url)\n",
    "def visualize_point_clouds(pc_A, vis=None):\n",
    "    if vis is None:\n",
    "        vis = meshcat.Visualizer(zmq_url)\n",
    "    vis[\"/Background\"].set_property('visible', False)\n",
    "    #vis[\"/Cameras/default/\"].set_transform(tf.translation_matrix([0, 0, 1]))\n",
    "    vis[\"/Cameras/default/rotated/<object>\"].set_property(\"zoom\", 10.5)\n",
    "\n",
    "    vis[\"red_bunny\"].set_object(g.PointCloud(pc_A, generate_color_mat([1, 0, 0], pc_A.shape), size=0.01))\n",
    "    return vis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Bva0aj0GlDSI"
   },
   "source": [
    "# Problem Description\n",
    "In the lecture, we learned about the RANSAC algorithm. In this exercise, you will implement the RANSAC algorithm to separate the Stanford bunny from its environment!\n",
    "\n",
    "**These are the main steps of the exercise:**\n",
    "1. Implement the `ransac` method.\n",
    "2. Implement the `remove_plane` method to remove the points that belong to the planar surface.\n",
    "\n",
    "Let's first visualize the point clouds of Stanford bunny in meshcat!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "XtPLYaZhlDSJ"
   },
   "outputs": [],
   "source": [
    "vis = visualize_point_clouds(bunny_w_plane, vis)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "zF54ocy-lDSR"
   },
   "source": [
    "You should notice that now there is a planar surface underneath the bunny. You may assume the bunny is currently placed on a table, where the planar surface is the tabletop. In this exercise, your objective is to remove the planar surface."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ki2f7sxZlDSS"
   },
   "source": [
    "A straightforward way to achieve a better fit is to remove the planar surface underneath the bunny. To do so, we provide you a function to fit a planar surface. \n",
    "\n",
    "Recall that a plane equation is of the form\n",
    "$$a x + b y + c z + d = 0$$\n",
    "where $[a,b,c]^T$ is a vector normal to the plane and (if it's a unit normal) $d$ is the negative of the distance from the origin to the plane in the direction of the normal.  We'll represent a plane by the vector $[a,b,c,d]$.\n",
    "\n",
    "The fitted planes are shown as translucent disks of radius $r$ centered at the points. The gray lines represent the planes' normals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "g2tMdp5PlDSS"
   },
   "outputs": [],
   "source": [
    "plane_equation = fit_plane(bunny_w_plane.T)\n",
    "print(plane_equation)\n",
    "DrawFacet(vis, plane_equation, 'naive_plane', center=[0,0,-plane_equation[-1]], thickness=0.005, radius=0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "HQs47D9SlDSX"
   },
   "source": [
    "You should notice that the planar surface cannot be fitted exactly either. This is because it takes account of all points in the scene to fit the plane. Since a significant portion of the point cloud belongs to the bunny, the fitted plane is noticeably elevated above the ground. \n",
    "\n",
    "To improve the result of the fitted plane, you will use RANSAC!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "8VB5_NrqlDSX"
   },
   "source": [
    "## RANSAC\n",
    "With the presence of outliers (bunny), we can use RANSAC to get more reliable estimates. The idea is to fit a plane using many random choices of a minimal set of points (3), fit a plane for each one, count how many points are within a distance tolerance to that plane (the inliers), and return the estimate with the most inliers.\n",
    "\n",
    "**Complete the function `ransac`.  It takes a data matrix, a tolerance, a value of iterations, and a model regressor. It returns an equation constructed by the model regressor and a count of inliers.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "8PZI0rU_lDSY"
   },
   "outputs": [],
   "source": [
    "def ransac(point_cloud, model_fit_func, tolerance=1e-3, max_iterations=500):\n",
    "    '''\n",
    "    Args:\n",
    "      point_cloud is (N, 3) numpy array\n",
    "      tolerance is a float\n",
    "      max_iterations is a (small) integer\n",
    "    Returns:\n",
    "      (4,) numpy array\n",
    "    '''\n",
    "    best_ic = 0                 # inlier count\n",
    "    best_model = np.ones(4)     # plane equation ((4,) array)\n",
    "\n",
    "    ##################\n",
    "    # your code here\n",
    "    ##################\n",
    "\n",
    "    return  best_ic, best_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ubTmMUnelDSi"
   },
   "source": [
    "Now you should have a lot better estimate of the planar surface with the use of RANSAC! Let's visualize the plane now!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "OST0NACZlDSi"
   },
   "outputs": [],
   "source": [
    "inlier_count, ransac_plane = ransac(bunny_w_plane.T, fit_plane, 0.001, 500)\n",
    "print(ransac_plane)\n",
    "DrawFacet(vis, ransac_plane, 'ransac_plane', center=[0,0,-ransac_plane[-1]], thickness=0.005, radius=0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "IjB0y8hAlDSm"
   },
   "source": [
    "## Remove Planar Surface\n",
    "\n",
    "Now all you need to do is to remove the points that belong to the planar surface. You may do so by rejecting all points that are \n",
    "\\begin{equation}\n",
    "|| a x + b y + c z + d || > tol\n",
    "\\end{equation}\n",
    "\n",
    "Note that since you are fitting a plane, the bunny is this case is the \"outlier\". Your job, however, is to keep the bunny and remove the planar surface.\n",
    "\n",
    "**Complete the function below to remove the points that belongs to the planar surface**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "M7Bt0qYqlDSn"
   },
   "outputs": [],
   "source": [
    "def remove_plane(point_cloud, ransac, tol=1e-4):\n",
    "    \"\"\"\n",
    "    Find the nearest (Euclidean) neighbor in point_cloud_B for each\n",
    "    point in point_cloud_A.\n",
    "    Args:\n",
    "        point_cloud: Nx3 numpy array of points\n",
    "        ransac: The RANSAC function to use (call ransac(args))\n",
    "        plane_equation: (4,) numpy array, contains the coefficients of the plane\n",
    "    Returns:\n",
    "        point_cloud_wo_plane: Nx3 numpy array of points\n",
    "    \"\"\"\n",
    "    point_cloud_wo_plane = np.zeros((100,3))\n",
    "    return point_cloud_wo_plane"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "oizYYZ1KlDSw"
   },
   "outputs": [],
   "source": [
    "bunny_wo_plane = remove_plane(bunny_w_plane.T, ransac)\n",
    "vis = visualize_point_clouds(bunny_wo_plane.T, vis)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "MwE8yNg58VQN"
   },
   "source": [
    "## How will this notebook be Graded?##\n",
    "\n",
    "If you are enrolled in the class, this notebook will be graded using [Gradescope](www.gradescope.com). You should have gotten the enrollement code on our announcement in Piazza. \n",
    "\n",
    "For submission of this assignment, you must do two things. \n",
    "- Download and submit the notebook `ransac.ipynb` to Gradescope's notebook submission section, along with your notebook for the other problems.\n",
    "- Copy and Paste your answer to the kinematic singularity problem to Gradescope's written submission section. \n",
    "\n",
    "We will evaluate the local functions in the notebook to see if the function behaves as we have expected. For this exercise, the rubric is as follows:\n",
    "- [4 pts] `ransac` must be implemented correctly. \n",
    "- [2 pts] `remove_plane` must be implemented correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "xj5nAh4g8VQO"
   },
   "outputs": [],
   "source": [
    "from manipulation.exercises.pose.test_ransac import TestRANSAC\n",
    "from manipulation.exercises.grader import Grader \n",
    "\n",
    "Grader.grade_output([TestRANSAC], [locals()], 'results.json')\n",
    "Grader.print_test_results('results.json')"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "colab": {
   "collapsed_sections": [],
   "name": "ransac.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}